#!/usr/bin/env python3
"""
简化版MCP WebSocket服务器

这是一个更稳定的MCP WebSocket服务器实现，专门用于解决连接问题。

使用方法：
python mcp_server_simple.py

作者: 张朝文
版本: 1.0.0
"""

import asyncio
import websockets
import json
import uuid
import signal
import sys
from datetime import datetime
from typing import Dict, Set
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('mcp_server_simple.log'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger('MCP_Simple_Server')

class SimpleMCPServer:
    """简化的MCP WebSocket服务器"""
    
    def __init__(self, host: str = 'localhost', port: int = 8080):
        self.host = host
        self.port = port
        self.clients: Set[websockets.WebSocketServerProtocol] = set()
        self.running = False
        
    async def handle_client(self, websocket, path=None):
        """处理客户端连接"""
        client_info = f"{websocket.remote_address[0]}:{websocket.remote_address[1]}"
        session_id = str(uuid.uuid4())
        
        # 获取请求路径
        if path is None:
            path = websocket.path if hasattr(websocket, 'path') else "/"
        
        logger.info(f"MCP客户端连接: {client_info}, 路径: {path}")
        
        # 检查路径是否支持 /mcp
        if path != "/mcp" and path != "/":
            logger.warning(f"不支持的路径: {path}")
            await websocket.close(code=1008, reason="不支持的路径")
            return
        
        self.clients.add(websocket)
        
        try:
            # 处理消息循环
            async for message in websocket:
                try:
                    data = json.loads(message)
                    message_type = data.get('type', 'UNKNOWN')
                    
                    logger.info(f"收到消息 [{session_id}]: {message_type}")
                    
                    if message_type == 'CONNECT':
                        # 发送连接确认
                        response = {
                            'type': 'CONNECT_ACK',
                            'session_id': session_id,
                            'status': 'connected',
                            'server_info': {
                                'name': 'Simple MCP Server',
                                'version': '1.0.0'
                            },
                            'timestamp': datetime.now().isoformat()
                        }
                        await websocket.send(json.dumps(response))
                        logger.info(f"发送CONNECT_ACK: {session_id}")
                        
                    elif message_type == 'REQUEST':
                        # 处理请求
                        message_id = data.get('message_id', str(uuid.uuid4()))
                        payload = data.get('payload', {})
                        query = payload.get('query', 'Hello')
                        
                        response_payload = {
                            'message': f"服务器收到请求: {query}",
                            'processed_at': datetime.now().isoformat(),
                            'request_id': message_id,
                            'server_response': f"已处理查询: {query}",
                            'status': 'success'
                        }
                        
                        response = {
                            'type': 'RESPONSE',
                            'message_id': message_id,  # 使用与请求相同的message_id
                            'payload': response_payload,
                            'status': 'success',
                            'timestamp': datetime.now().isoformat()
                        }
                        
                        await websocket.send(json.dumps(response))
                        logger.info(f"发送RESPONSE: {message_id}")
                        
                    elif message_type == 'PING':
                        # 处理心跳
                        response = {
                            'type': 'PONG',
                            'timestamp': datetime.now().isoformat(),
                            'server_time': datetime.now().isoformat()
                        }
                        await websocket.send(json.dumps(response))
                        logger.debug(f"发送PONG: {session_id}")
                        
                    else:
                        # 未知消息类型
                        error_response = {
                            'type': 'ERROR',
                            'error_code': 'UNKNOWN_MESSAGE_TYPE',
                            'error_message': f'未知的消息类型: {message_type}',
                            'timestamp': datetime.now().isoformat()
                        }
                        await websocket.send(json.dumps(error_response))
                        logger.warning(f"未知消息类型: {message_type}")
                        
                except json.JSONDecodeError as e:
                    logger.error(f"JSON解析错误: {str(e)}")
                    error_response = {
                        'type': 'ERROR',
                        'error_code': 'INVALID_JSON',
                        'error_message': '无效的JSON格式',
                        'timestamp': datetime.now().isoformat()
                    }
                    await websocket.send(json.dumps(error_response))
                    
                except Exception as e:
                    logger.error(f"处理消息错误: {str(e)}")
                    error_response = {
                        'type': 'ERROR',
                        'error_code': 'PROCESSING_ERROR',
                        'error_message': f'处理消息时发生错误: {str(e)}',
                        'timestamp': datetime.now().isoformat()
                    }
                    await websocket.send(json.dumps(error_response))
                    
        except websockets.exceptions.ConnectionClosed:
            logger.info(f"MCP客户端断开连接: {client_info}")
        except websockets.exceptions.InvalidMessage as e:
            logger.warning(f"无效的WebSocket消息 [{client_info}]: {str(e)}")
        except websockets.exceptions.ConnectionClosedError as e:
            logger.info(f"WebSocket连接关闭 [{client_info}]: {str(e)}")
        except Exception as e:
            logger.error(f"客户端连接错误 [{client_info}]: {str(e)}")
        finally:
            # 清理连接
            self.clients.discard(websocket)
            logger.info(f"客户端连接清理完成: {client_info}")
    
    async def start_server(self):
        """启动MCP服务器"""
        try:
            logger.info(f"启动简化版MCP服务器: ws://{self.host}:{self.port}")
            
            # 启动WebSocket服务器，添加更好的错误处理
            async with websockets.serve(
                self.handle_client,
                self.host,
                self.port,
                subprotocols=None,
                ping_interval=20,
                ping_timeout=10,
                close_timeout=10,
                max_size=2**20,  # 1MB
                max_queue=2**5,  # 32
                compression=None
            ) as server:
                
                self.running = True
                logger.info(f"MCP服务器已启动，监听地址: ws://{self.host}:{self.port}")
                logger.info("支持路径: /mcp, /")
                logger.info("等待客户端连接...")
                
                # 等待服务器关闭
                await server.wait_closed()
                
        except Exception as e:
            logger.error(f"MCP服务器启动失败: {str(e)}")
            raise
    
    def get_server_status(self):
        """获取服务器状态"""
        return {
            'running': self.running,
            'host': self.host,
            'port': self.port,
            'client_count': len(self.clients)
        }


# 全局服务器实例
mcp_server = None


def signal_handler(signum, frame):
    """信号处理器"""
    logger.info(f"收到信号 {signum}，正在停止服务器...")
    if mcp_server:
        mcp_server.running = False
    sys.exit(0)


async def main():
    """主函数"""
    global mcp_server
    
    # 注册信号处理器
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    logger.info("=" * 60)
    logger.info("简化版MCP WebSocket服务器启动")
    logger.info("=" * 60)
    
    try:
        # 创建并启动MCP服务器
        mcp_server = SimpleMCPServer(host='localhost', port=8080)
        await mcp_server.start_server()
        
    except KeyboardInterrupt:
        logger.info("收到键盘中断信号")
    except Exception as e:
        logger.error(f"服务器运行错误: {str(e)}")
    finally:
        logger.info("MCP服务器已停止")


if __name__ == "__main__":
    """启动简化版MCP服务器"""
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("服务器已停止")
    except Exception as e:
        logger.error(f"启动失败: {str(e)}")
        sys.exit(1)
